# include "scene.h"
# include "discr.h"
# include <stdio.h>
# include <math.h>
# include <stdlib.h>
/*
# include <time.h>
# include <sys/timeb.h>*/

# ifdef _OMP_VEL_ADVANCE_
# include <omp.h>
# endif
/*
# ifdef CALL_SOLVE_POISSON_OMP
extern long int solve_poisson_omp(struct grid *g, const double ***d, unsigned short iRK, double ***fspace);
# else*/
extern long int solve_poisson(struct grid *g, const double ***d, unsigned short iRK);
// # endif

extern void vel_bc(struct grid *g, short tflag, unsigned short iRK);
extern void press_bc(const struct grid *g, short tflag, unsigned short iRK);

extern double dens(double T);
extern double kvisc(double T);

static double getA(int i, int j, int k, const struct grid *g);
static double getB(int i, int j, int k, const struct grid *g);
static double getC(int i, int j, int k, const struct grid *g);


struct diverg_max
     {
      short im, jm, km;
      double divmax;
     };

static void divergence(const struct grid *g, double ***diverg, struct diverg_max *mdiv);

extern unsigned int vel_advance(struct grid *g, unsigned int *pcnt_max, const struct arr_struct *depVar, const struct arr_struct *memFace)
{
 unsigned int pcnt, pcnt_subttl;
 short i, j, k, _i, _j, _k, nr, nth, nz;

 double density0, *r, *th, *z, ***d, ***D0, ***p, ***_u[nRK+1], ***_v[nRK+1], ***_w[nRK+1], ***_p[nRK],\
        ***_A[nRK], ***_B[nRK], ***_C[nRK], ***divH0, ***divH1,/* divH2,*/ ***div2p0,/* div2p1,*/ dt;

// struct timeb t1, t2;

#ifdef TRACE_FILE
 fprintf(g->trptr, "Entering into 'vel_advance()'...\n");
#endif

 r = g->r; th = g->th; z = g->z; dt = g->dt;
 nr = g->nr; nth = g->nth; nz = g->nz;
 *pcnt_max = pcnt_subttl = 0;
 density0 = g->density0;

  _u[0] = g->u;            _v[0] = g->v;             _w[0] = g->w;            _p[0] = g->p;
  _u[1] = depVar->q[4];    _v[1] = depVar->q[5];     _w[1] = depVar->q[6];    _p[1] = depVar->q[7];
  _u[2] = depVar->q[8];    _v[2] = depVar->q[9];     _w[2] = depVar->q[10];   _p[2] = depVar->q[11];
  _u[3] = depVar->q[12];   _v[3] = depVar->q[13];    _w[3] = depVar->q[14];
  D0 = depVar->q[11];      divH0 = depVar->q[12];    divH1 = depVar->q[13];   div2p0 = depVar->q[14];

  for(i = 0; i < nRK; i++)
    {
     _A[i] = memFace[0].q[i];
     _B[i] = memFace[1].q[i];
     _C[i] = memFace[2].q[i];
    }

/* Calculating divergence field at the current n-th time-step. It gets filled in D0[][][]. */
  divergence((const struct grid *)g, D0, (struct diverg_max *)0);

#ifdef _OMP_VEL_ADVANCE_
 #pragma omp parallel private(i,j,k,_i,_j,_k)
#endif
 {
/* --------------------------- Calculation of u(1), v(1), w(1), and p(1) ---------------------------------- */
// ftime(&t1);
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp for schedule(guided)
# endif
  for(k = 0; k <= nz; k++) for(i = 0; i <= nr; i++) for(j = 0; j <= nth; j++)
         {
          if(j < nth && k < nz) _A[0][i][j][k] = getA(NrG+i-1, NthG+j, NzG+k, (const struct grid *)g);
          if(i < nr && k < nz)  _B[0][i][j][k] = getB(NrG+i, NthG+j-1, NzG+k, (const struct grid *)g);
          if(i < nr && j < nth) _C[0][i][j][k] = getC(NrG+i, NthG+j, NzG+k-1, (const struct grid *)g);
         }

# ifdef _OMP_VEL_ADVANCE_
  #pragma omp for schedule(guided)
# endif
  for(_k = 0; _k < nz; _k++) for(k = _k+NzG, _i = 0; _i < nr; _i++) for(i = _i+NrG, _j = 0, j = NthG; _j < nth; j++, _j++)
         {
          _u[1][i][j][k] = (i == NrG+nr-1 ? 0:_u[0][i][j][k]\
              + c2*dt*( _A[0][_i+1][_j][_k] - (2.0/density0)*(_p[0][i+1][j][k] - _p[0][i][j][k])/(r[i+2]-r[i]) ) );

          _v[1][i][j][k] = _v[0][i][j][k]\
             + c2*dt*( _B[0][_i][_j+1][_k] - (4.0/density0)*(_p[0][i][j+1][k] - _p[0][i][j][k])/((r[i]+r[i+1])*(th[j+2]-th[j])) );

          _w[1][i][j][k] = (k == NzG+nz-1 ? 0:_w[0][i][j][k]\
             + c2*dt*( _C[0][_i][_j][_k+1] - (2.0/density0)*(_p[0][i][j][k+1] - _p[0][i][j][k])/(z[k+2]-z[k]) ) );
         }

# ifdef _OMP_VEL_ADVANCE_
  #pragma omp single
# endif
    {
     g->u = _u[1]; g->v = _v[1]; g->w = _w[1];
     vel_bc(g, 0, 1);
    }


/* _u[1], _v[1], _w[1] are used in "press_bc()". Therefore, the memory used by these can't be distroyed. */
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp for schedule(guided)
# endif
  for(k = 0; k <= nz; k++) for(i = 0; i <= nr; i++) for(j = 0; j <= nth; j++)
         {
          if(j < nth && k < nz) _A[1][i][j][k] = getA(NrG+i-1, NthG+j, NzG+k, (const struct grid *)g);
          if(i < nr && k < nz)  _B[1][i][j][k] = getB(NrG+i, NthG+j-1, NzG+k, (const struct grid *)g);
          if(i < nr && j < nth) _C[1][i][j][k] = getC(NrG+i, NthG+j, NzG+k-1, (const struct grid *)g);
         }

/* Calculating RHS of the PPE-1. */
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp single
# endif
   {
    p = _p[0]; d = depVar->q[8];
   }

# ifdef _OMP_VEL_ADVANCE_
  #pragma omp for schedule(guided)
# endif
  for(_k = 0; _k < nz; _k++) for(k = _k+NzG, _i = 0; _i < nr; _i++) for(i = _i+NrG, _j = 0, j = NthG; _j < nth; j++, _j++)
      {
       div2p0[i][j][k] = DIV2P_ijk;

       divH0[i][j][k] = 2*(( r[i+1]*_A[0][_i+1][_j][_k] - (_i == 0 ? _A[0][_i][_j][_k]:r[i]*_A[0][_i][_j][_k]))/(r[i+1]-r[i])\
                           + (_B[0][_i][_j+1][_k] - _B[0][_i][_j][_k])/(th[j+1]-th[j]) )/(r[i]+r[i+1])\
                           + (_C[0][_i][_j][_k+1] - _C[0][_i][_j][_k])/(z[k+1]-z[k]);

       divH1[i][j][k] = 2*( (r[i+1]*_A[1][_i+1][_j][_k] - (_i == 0 ? _A[1][_i][_j][_k]:r[i]*_A[1][_i][_j][_k]))/(r[i+1]-r[i])\
                          + (_B[1][_i][_j+1][_k] - _B[1][_i][_j][_k])/(th[j+1]-th[j]) )/(r[i]+r[i+1])\
                          + (_C[1][_i][_j][_k+1] - _C[1][_i][_j][_k])/(z[k+1]-z[k]);

       d[i][j][k] = density0*(divH1[i][j][k] + (D0[i][j][k]/dt + a31*(divH0[i][j][k] - div2p0[i][j][k]/density0))/a32);
      }


/* Solving the PPE(1) to calculate p(1). */

/* Copying _p[0] into _p[1] to provide it as initial condition for pressure while solving the PPE. */
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp for schedule(static)
# endif
  for(i = 0; i < 2*NrG+nr; i++) for(j = 0; j < 2*NthG+nth; j++) for(k = 0; k < 2*NzG+nz; k++)
  _p[1][i][j][k] = _p[0][i][j][k];


# ifdef _OMP_VEL_ADVANCE_
  #pragma omp single
# endif
   {
    fprintf(g->dptr, "  [ PPE - 1 ]\n");
    fprintf(g->tptr, "PPE-1:-[( "); fflush(g->tptr);
    g->p = _p[1];
    pcnt = solve_poisson(g, (const double ***)d, 1);
    fprintf(g->tptr, ")]; ");  fflush(g->tptr);

    pcnt_subttl += pcnt;
    if(pcnt > *pcnt_max) *pcnt_max = pcnt;
   }
/* ------------------------------------------------------------------------------------------------------------------------- */


/* ------------------------------------- Calculation of u(2), v(2), w(2), and p(2) ----------------------------------------- */
/* Calculating u(2), v(2), and w(2). */
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp for schedule(guided)
# endif
  for(_k = 0; _k < nz; _k++) for(k = _k+NzG, _i = 0; _i < nr; _i++) for(i = _i+NrG, _j = 0, j = NthG; _j < nth; j++, _j++)
         {
          _u[2][i][j][k] = (i == NrG+nr-1 ? 0:_u[0][i][j][k]\
                + a31*dt*( _A[0][_i+1][_j][_k] - (2.0/density0)*(_p[0][i+1][j][k] - _p[0][i][j][k])/(r[i+2] - r[i]) )\
                + a32*dt*( _A[1][_i+1][_j][_k] - (2.0/density0)*(_p[1][i+1][j][k] - _p[1][i][j][k])/(r[i+2] - r[i]) ) );

          _v[2][i][j][k] = _v[0][i][j][k]\
                + a31*dt*( _B[0][_i][_j+1][_k] - (4.0/density0)*(_p[0][i][j+1][k] - _p[0][i][j][k])/( (r[i]+r[i+1])*(th[j+2]-th[j]) ) )\
                + a32*dt*( _B[1][_i][_j+1][_k] - (4.0/density0)*(_p[1][i][j+1][k] - _p[1][i][j][k])/( (r[i]+r[i+1])*(th[j+2]-th[j]) ) );

          _w[2][i][j][k] = (k == NzG+nz-1 ? 0:_w[0][i][j][k]\
                               + a31*dt*( _C[0][_i][_j][_k+1] - (2.0/density0)*(_p[0][i][j][k+1] - _p[0][i][j][k])/(z[k+2]-z[k]) )\
                               + a32*dt*( _C[1][_i][_j][_k+1] - (2.0/density0)*(_p[1][i][j][k+1] - _p[1][i][j][k])/(z[k+2]-z[k]) ) );
         }

# ifdef _OMP_VEL_ADVANCE_
  #pragma omp single
# endif
   {
    g->u = _u[2]; g->v = _v[2]; g->w = _w[2];
    vel_bc(g, 0, 2);
   }

/* _u[2], _v[2], _w[2] are used in "press_bc()". Therefore, the memory used by these can't be distroyed. */
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp for schedule(guided)
# endif
  for(k = 0; k <= nz; k++) for(i = 0; i <= nr; i++) for(j = 0; j <= nth; j++)
         {
          if(j < nth && k < nz) _A[2][i][j][k] = getA(NrG+i-1, NthG+j, NzG+k, (const struct grid *)g);
          if(i < nr && k < nz)  _B[2][i][j][k] = getB(NrG+i, NthG+j-1, NzG+k, (const struct grid *)g);
          if(i < nr && j < nth) _C[2][i][j][k] = getC(NrG+i, NthG+j, NzG+k-1, (const struct grid *)g);
         }

/* Calculation of R.H.S. of the PPE(2). */
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp single
# endif
   {
    p = _p[1]; d = depVar->q[4];
   }

# ifdef _OMP_VEL_ADVANCE_
  #pragma omp for schedule(guided)
# endif
  for(_k = 0; _k < nz; _k++) for(k = _k+NzG, _i = 0; _i < nr; _i++) for(i = _i+NrG, _j = 0, j = NthG; _j < nth; j++, _j++)
      {
#      define divH2 (2*( (r[i+1]*_A[2][_i+1][_j][_k] - (_i == 0 ? _A[2][_i][_j][_k]:r[i]*_A[2][_i][_j][_k]))/(r[i+1]-r[i])\
                    + (_B[2][_i][_j+1][_k] - _B[2][_i][_j][_k])/(th[j+1]-th[j]) )/(r[i]+r[i+1])\
                    + (_C[2][_i][_j][_k+1] - _C[2][_i][_j][_k])/(z[k+1]-z[k]))

       d[i][j][k] = density0*( divH2 + ( D0[i][j][k]/dt + a41*(divH0[i][j][k] - div2p0[i][j][k]/density0)\
                                                       + a42*(divH1[i][j][k] - DIV2P_ijk/density0) )/a43 );
#      undef divH2
      }

/* Solving the PPE(2) to calculate p(2). */

/* Copying _p[1] into _p[2] to provide it as initial condition for pressure while solving the PPE. */
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp for schedule(static)
# endif
  for(i = 0; i < 2*NrG+nr; i++) for(j = 0; j < 2*NthG+nth; j++) for(k = 0; k < 2*NzG+nz; k++)
      _p[2][i][j][k] = _p[1][i][j][k];


# ifdef _OMP_VEL_ADVANCE_
  #pragma omp single
# endif
   {
    fprintf(g->dptr, "\n  [ PPE - 2 ]\n");
    fprintf(g->tptr, "PPE-2:-[( "); fflush(g->tptr);

    g->p = _p[2];

    pcnt = solve_poisson(g, (const double ***)d, 2);
    fprintf(g->tptr, ")]; ");  fflush(g->tptr);

    pcnt_subttl += pcnt;
    if(pcnt > *pcnt_max) *pcnt_max = pcnt;
   }
/* ---------------------------------------------------------------------------------------------------------- */

/* ================================= Calculation of u, v, w, and p at next time-step ========================== */
/* Calculation of u, v, and w at next time-step. */
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp for schedule(guided)
# endif
  for(_k = 0; _k < nz; _k++) for(k = _k+NzG, _i = 0; _i < nr; _i++) for(i = _i+NrG, _j = 0, j = NthG; _j < nth; j++, _j++)
         {
          _u[3][i][j][k] = (i == NrG+nr-1 ? 0:_u[0][i][j][k]\
                  + b1*dt*( _A[0][_i+1][_j][_k] - (2.0/density0)*(_p[0][i+1][j][k] - _p[0][i][j][k])/(r[i+2]-r[i]) )\
                  + b2*dt*( _A[1][_i+1][_j][_k] - (2.0/density0)*(_p[1][i+1][j][k] - _p[1][i][j][k])/(r[i+2]-r[i]) )\
                  + b3*dt*( _A[2][_i+1][_j][_k] - (2.0/density0)*(_p[2][i+1][j][k] - _p[2][i][j][k])/(r[i+2]-r[i]) ) );

          _v[3][i][j][k] = _v[0][i][j][k]\
                  + b1*dt*( _B[0][_i][_j+1][_k] - (4.0/density0)*(_p[0][i][j+1][k] - _p[0][i][j][k])/((r[i]+r[i+1])*(th[j+2]-th[j])) )\
                  + b2*dt*( _B[1][_i][_j+1][_k] - (4.0/density0)*(_p[1][i][j+1][k] - _p[1][i][j][k])/((r[i]+r[i+1])*(th[j+2]-th[j])) )\
                  + b3*dt*( _B[2][_i][_j+1][_k] - (4.0/density0)*(_p[2][i][j+1][k] - _p[2][i][j][k])/((r[i]+r[i+1])*(th[j+2]-th[j])) );

          _w[3][i][j][k] = (k == NzG+nz-1 ? 0:_w[0][i][j][k]\
                  + b1*dt*( _C[0][_i][_j][_k+1] - (2.0/density0)*(_p[0][i][j][k+1] - _p[0][i][j][k])/(z[k+2]-z[k]) )\
                  + b2*dt*( _C[1][_i][_j][_k+1] - (2.0/density0)*(_p[1][i][j][k+1] - _p[1][i][j][k])/(z[k+2]-z[k]) )\
                  + b3*dt*( _C[2][_i][_j][_k+1] - (2.0/density0)*(_p[2][i][j][k+1] - _p[2][i][j][k])/(z[k+2]-z[k]) ) );
         }

# ifdef _OMP_VEL_ADVANCE_
  #pragma omp single
# endif
   {
    g->u = _u[3]; g->v = _v[3]; g->w = _w[3];
    vel_bc(g, 1, 0);

    g->um = depVar->q[4]; g->vm = depVar->q[5]; g->wm = depVar->q[6]; //To store mean velocities [i.e. velocity at time = t+1/2.]
   }

/* Storing mean velocities and copying the newly evaluated velocities to _u[0], _v[0], and _w[0]. */
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp for schedule(static) nowait
# endif
  for(k = 0; k < nz+2*NzG; k++) for(i = 0; i < nr+2*NrG; i++) for(j = 0; j < nth+2*NthG; j++)
      {
    /* Storing mean-velocities. */
       g->um[i][j][k] = 0.5*(_u[0][i][j][k] + _u[3][i][j][k]);
       g->vm[i][j][k] = 0.5*(_v[0][i][j][k] + _v[3][i][j][k]);
       g->wm[i][j][k] = 0.5*(_w[0][i][j][k] + _w[3][i][j][k]);

    /* Copying newly evaluated velocities. */
       _u[0][i][j][k] = _u[3][i][j][k];
       _v[0][i][j][k] = _v[3][i][j][k];
       _w[0][i][j][k] = _w[3][i][j][k];
      }

# ifdef _OMP_VEL_ADVANCE_
  #pragma omp for schedule(static)
# endif
  for(k = NzG; k < nz + NzG; k++) for(i = NrG; i < nr + NrG; i++) for(j = NthG; j < nth + NthG; j++)
   /* Pressure at the next time-step will be given by _p[0]. */
      _p[0][i][j][k] = d1*_p[0][i][j][k] + d2*_p[1][i][j][k] + d3*_p[2][i][j][k];
 }


  g->u = _u[0]; g->v = _v[0]; g->w = _w[0];  g->p = _p[0];
  press_bc(g, 1, 0);

/* ---------------------------------------------------------------------------------------------------------------- */

#ifdef TRACE_FILE
 fprintf(g->trptr, "...'vel_advance()' ends.\n");
#endif
 return pcnt_subttl;

/* q[0]~q[6] are engagged. These can't be used in other functions. */
}



static double getA(int i, int j, int k, const struct grid *g)
 {
  double ***u, ***v, ***w, ***T, *r, *th, *z, kvisc0;

  u = g->u;  v = g->v;  w = g->w; T = g->T; r = g->r;  th = g->th;  z = g->z;
  kvisc0 = g->kvisc0;

  return (i == NrG-1) ? rA_CENTER : AA;
 }


static double getB(int i, int j, int k, const struct grid *g)
 {
  double ***u, ***v, ***w, ***T, *r, *th, *z, kvisc0;

  u = g->u;  v = g->v;  w = g->w; T = g->T;  r = g->r;  th = g->th;  z = g->z;
  kvisc0 = g->kvisc0;

  return BB;
 }


static double getC(int i, int j, int k, const struct grid *g)
 {
  double ***u, ***v, ***w, ***T, *r, *th, *z, kvisc0, density0;

  r = g->r;  th = g->th;  z = g->z; u = g->u;  v = g->v;  w = g->w; T = g->T;
  kvisc0 = g->kvisc0; density0 = g->density0;

  return CC;
 }


static void divergence(const struct grid *g, double ***diverg, struct diverg_max *mdiv)
   {
    double *r, *th, *z, ***u, ***v, ***w, *div;
    short i, j, k, tid, nTRD;
    struct diverg_max *tddiv;

#   ifdef TRACE_FILE
    fprintf(g->trptr, "Entering into 'divergence()'...\n");
#   endif

    u = g->u;  v = g->v;  w = g->w; r = g->r; th = g->th; z = g->z;

    nTRD = 1;
#   ifdef _OMP_VEL_ADVANCE_
    nTRD = g->nTRD;
#   endif

    tddiv = (struct diverg_max *)malloc(sizeof(struct diverg_max)*nTRD);
    div = (double *)malloc(sizeof(double)*nTRD);

  /* Evaluating divergence: */
    tid = 0;
#   ifdef _OMP_VEL_ADVANCE_
    #pragma omp parallel private(i,j,k,tid)
#   endif
         {
#         ifdef _OMP_VEL_ADVANCE_
          tid = omp_get_thread_num();
#         endif

          tddiv[tid].im = -1;
          tddiv[tid].jm = -1;
          tddiv[tid].km = -1;
          tddiv[tid].divmax = 0.0;

#         ifdef _OMP_VEL_ADVANCE_
          #pragma omp for schedule(guided)
#         endif
          for(i = NrG; i < g->nr+NrG; i++) for(j = NthG; j < g->nth+NthG; j++) for(k = NzG; k < g->nz+NzG; k++)
             {
              div[tid] = DIV_ijk;

           /* Filtering out the maximum divergence. */
              if((div[tid] < 0 ? -div[tid]:div[tid]) > (tddiv[tid].divmax < 0 ? -tddiv[tid].divmax:tddiv[tid].divmax))
                 {
                  tddiv[tid].divmax = div[tid];
                  tddiv[tid].im = i;
                  tddiv[tid].jm = j;
                  tddiv[tid].km = k;
                 }

              diverg[i][j][k] = div[tid];
             }
         }

    free(div);

 /* Calculating the absolute maximum value of the divergence. */
    for(i = 0, tid = 1; tid < nTRD; tid++)
    if((tddiv[tid].divmax < 0 ? -tddiv[tid].divmax:tddiv[tid].divmax) > (tddiv[i].divmax < 0 ? -tddiv[i].divmax:tddiv[i].divmax)) i = tid;

    fprintf(g->tptr, "divmax(%d, %d, %d) = %E; ", tddiv[i].im, tddiv[i].jm, tddiv[i].km, tddiv[i].divmax); fflush(g->tptr);

    if( mdiv != (struct diverg_max *)0 )
      { mdiv->divmax = tddiv[i].divmax; mdiv->im = tddiv[i].im; mdiv->jm = tddiv[i].jm; mdiv->km = tddiv[i].km; }

    free(tddiv);

#   ifdef TRACE_FILE
    fprintf(g->trptr, "...'divergence()' ends.\n");
#   endif

    return;
   }
